import psycopg2
import os
from openai import OpenAI
from pgvector.psycopg2 import register_vector
from typing import List, Dict, Any
import numpy as np

class VectorSearchEngine:
    def __init__(self, db_config, embedding_api_key=None, generation_api_key=None):
        """
        初始化向量搜索引擎
        :param db_config: 数据库连接配置字典
        :param embedding_api_key: 阿里云向量嵌入服务的API密钥
        :param generation_api_key: DeepSeek生成服务的API密钥
        """
        self.db_config = db_config
        
        # 创建两个不同的客户端，使用不同的API密钥
        self.embedding_client = OpenAI(
            api_key=embedding_api_key,
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
        )
        self.generation_client = OpenAI(
            api_key=generation_api_key,
            base_url="https://api.deepseek.com/v1"
        )
        self._initialize_database()

    def _initialize_database(self):
        """初始化数据库连接并注册向量类型"""
        self.conn = psycopg2.connect(**self.db_config)
        register_vector(self.conn)
        self._create_index_if_not_exists()

    def _create_index_if_not_exists(self):
        """创建HNSW索引（如果不存在）"""
        create_index_query = """
            CREATE INDEX IF NOT EXISTS collection_name
            ON vba_docs
            USING hnsw (vector vector_cosine_ops)
            WITH (m = 16, ef_construction = 64);
        """
        with self.conn.cursor() as cursor:
            cursor.execute(create_index_query)
        self.conn.commit()

    def get_embedding_from_api(self, text):
        """从API获取文本嵌入向量"""
        try:
            completion = self.embedding_client.embeddings.create(
                model="text-embedding-v3",
                input=text,
                dimensions=768,
                encoding_format="float"
            )
            # 从响应中提取向量
            embedding = completion.data[0].embedding
            return embedding
        except Exception as e:
            print(f"获取嵌入向量时出错: {str(e)}")
            return None

    def generate_from_api(self, prompt):
        """从API获取生成的文本"""
        try:
            response = self.generation_client.chat.completions.create(
                model="deepseek-chat",
                messages=[
                    {"role": "system", "content": "你是一个专业的VBA助手，请基于给定的上下文信息回答用户的问题。如果无法从上下文中找到相关信息，请明确告知。回答要简洁、准确，并尽可能引用相关的具体内容。"},
                    {"role": "user", "content": prompt}
                ],
                stream=False
            )
            return response.choices[0].message.content
        except Exception as e:
            print(f"生成文本时出错: {str(e)}")
            return None

    def generate_query_variations(self, query_text: str) -> List[str]:
        """
        生成查询的多个变体
        :param query_text: 原始查询文本
        :return: 查询变体列表
        """
        try:
            prompt = f"""请为以下查询生成3个不同角度的变体，每个变体都应该保持原始查询的核心含义，但从不同的视角表达。
            原始查询：{query_text}
            
            请直接返回3个变体，每行一个，不要添加任何额外的说明或编号。"""

            response = self.generation_client.chat.completions.create(
                model="deepseek-chat",
                messages=[
                    {"role": "system", "content": "你是一个专业的查询重写助手，请生成准确、多样的查询变体。"},
                    {"role": "user", "content": prompt}
                ],
                stream=False
            )
            
            variations = response.choices[0].message.content.strip().split('\n')
            # 确保至少返回原始查询
            variations = [v.strip() for v in variations if v.strip()]
            variations.append(query_text)
            return list(set(variations))  # 去重
        except Exception as e:
            print(f"生成查询变体时出错: {str(e)}")
            return [query_text]  # 如果生成失败，返回原始查询

    def get_embeddings_batch(self, texts: List[str]) -> List[List[float]]:
        """
        批量获取文本的嵌入向量
        :param texts: 文本列表
        :return: 向量列表
        """
        try:
            completion = self.embedding_client.embeddings.create(
                model="text-embedding-v3",
                input=texts,
                dimensions=768,
                encoding_format="float"
            )
            return [item.embedding for item in completion.data]
        except Exception as e:
            print(f"批量获取嵌入向量时出错: {str(e)}")
            return []

    def multi_query_search(self, query_text: str, top_n: int = 5, ef_search: int = 100, 
                          similarity_threshold: float = 0.5) -> List[Dict[str, Any]]:
        """
        执行多视角查询搜索
        :param query_text: 查询文本
        :param top_n: 每个查询返回的结果数量
        :param ef_search: HNSW搜索参数
        :param similarity_threshold: 相似度阈值
        :return: 合并后的检索结果
        """
        # 生成查询变体
        query_variations = self.generate_query_variations(query_text)
        
        # 批量获取查询变体的向量
        query_embeddings = self.get_embeddings_batch(query_variations)
        if not query_embeddings:
            return []

        all_results = []
        seen_ids = set()  # 用于去重

        try:
            with self.conn.cursor() as cursor:
                cursor.execute(f"SET hnsw.ef_search = {ef_search};")
                
                # 对每个查询变体执行搜索
                for embedding in query_embeddings:
                    query = """
                        SELECT id, content, metadata, 1 - (vector <=> %s::vector) AS similarity
                        FROM vba_docs
                        WHERE 1 - (vector <=> %s::vector) > %s
                        ORDER BY similarity DESC
                        LIMIT %s;
                    """
                    cursor.execute(query, (embedding, embedding, similarity_threshold, top_n))
                    results = cursor.fetchall()
                    
                    # 添加未见过的结果
                    for row in results:
                        if row[0] not in seen_ids:  # 使用ID去重
                            seen_ids.add(row[0])
                            all_results.append({
                                "id": row[0],
                                "content": row[1],
                                "source": row[2],
                                "similarity": float(row[3])
                            })

        except Exception as e:
            self.conn.rollback()
            print(f"多视角查询搜索时出错: {str(e)}")
            return []

        # 按相似度排序并返回前top_n个结果
        all_results.sort(key=lambda x: x["similarity"], reverse=True)
        return all_results[:top_n]

    def similarity_search(self, query_text, top_n=5, ef_search=100, generate_answer=True, similarity_threshold=0.5):
        """
        相似度查询并生成回答
        :param query_text: 查询文本
        :param top_n: 返回结果数量
        :param ef_search: HNSW搜索参数
        :param generate_answer: 是否生成回答
        :param similarity_threshold: 相似度阈值
        :return: 检索结果或生成的回答
        """
        # 使用多视角查询替代单一查询
        results = self.multi_query_search(query_text, top_n, ef_search, similarity_threshold)
        
        if not results:
            return {"answer": "抱歉，我找不到相关的信息来回答这个问题。", "context": "", "sources": []}

        contexts = []
        sources = []
        similarity_scores = []
        
        for result in results:
            contexts.append(result["content"])
            sources.append(result["source"])
            similarity_scores.append(result["similarity"])

        if generate_answer:
            context_text = "\n".join(contexts)
            answer = self.generate_with_context(query_text, context_text)
            return {
                "answer": answer,
                "context": context_text,
                "sources": sources,
                "similarity_scores": similarity_scores
            }
        else:
            return results

    def generate_with_context(self, query_text, context_text):
        """
        使用检索到的上下文生成回答
        :param query_text: 用户查询文本
        :param context_text: 检索到的上下文文本
        :return: 生成的回答
        """
        try:
            prompt = f"""请基于以下参考信息回答用户的问题。如果无法从参考信息中找到相关内容，请明确说明。
            
            参考信息：
            {context_text}

            用户问题：{query_text}

            请提供准确、简洁的回答："""

            answer = self.generate_from_api(prompt)
            return answer if answer else "抱歉，生成回答时出现错误。"

        except Exception as e:
            return f"生成回答时出错: {str(e)}"

    def __del__(self):
        """清理资源"""
        if hasattr(self, 'conn'):
            self.conn.close()